/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
 * Description: op remote look up execute.
 * Author: qingjicheng
 * Create: 2022-11-28
 */

#include <algorithm>

#include "transport_heterog_pub.h"
#include "transport_roce_pub.h"
#include "transport_shm_event_pub.h"
#include "log.h"
#include "adapter_hdds.h"
#include "adapter_rts.h"
#include "profiling_manager.h"
#include "adapter_prof.h"
#include "op_remote_look_up_executor.h"
#include "transport_heterog_event_roce_pub.h"
#include "remote_notify.h"
#include "ffts_common_pub.h"

using namespace std;

namespace hccl {

constexpr s32 PIPELINE_SEND_KEYS_SUB_STREAM_INDEX = 0;
constexpr s32 PIPELINE_RECV_VALUES_SUB_STREAM_INDEX = 1;
constexpr s32 PIPELINE_RECOVER_VALUE_SUB_STREAM_INDEX = 2;
constexpr s32 PIPELINE_SEND_KEYS_NOTIFY_INDEX = 1;
constexpr s32 PIPELINE_RECV_VALUES_NOTIFY_INDEX = 2;
constexpr s32 PIPELINE_RECOVER_VALUE_NOTIFY_INDEX = 3;
constexpr u32 NOTIFY_WAIT_TIME = 10;
constexpr u32 NOTIFY_RESET_UNIQUE_REVERSE_IDX = 2;

OpRemoteLookUpExecutor::OpRemoteLookUpExecutor(const std::unique_ptr<MrManager> &mrManager,
    const HcclDispatcher dispatcher, const u32 localWorldRank) : taskExeFuncTable_({
    {OpTaskType::SEND, std::bind(&OpRemoteLookUpExecutor::TaskSend, this, std::placeholders::_1)},
    {OpTaskType::RECV, std::bind(&OpRemoteLookUpExecutor::TaskRecv, this, std::placeholders::_1)},
    {OpTaskType::MEMCPY, std::bind(&OpRemoteLookUpExecutor::TaskMemcpy, this, std::placeholders::_1)},
    {OpTaskType::DROP_DUPLICATES, std::bind(&OpRemoteLookUpExecutor::TaskDropDuplicates, this, std::placeholders::_1)},
    {OpTaskType::DROP_DUPLICATES_KERNEL_LAUNCH,
        std::bind(&OpRemoteLookUpExecutor::TaskKeyDropDuplicatesKernelLaunch, this, std::placeholders::_1)},
    {OpTaskType::SEND_KEYS_KERNEL_LAUNCH,
        std::bind(&OpRemoteLookUpExecutor::TaskSendKeysKernelLaunch, this, std::placeholders::_1)},
    {OpTaskType::EVENT_SEND_RECORD,
        std::bind(&OpRemoteLookUpExecutor::TaskSendKeysRecord, this, std::placeholders::_1)},
    {OpTaskType::EVENT_WAIT_RECV_DONE,
        std::bind(&OpRemoteLookUpExecutor::TaskEventWaitRecvDone, this, std::placeholders::_1)},
    {OpTaskType::RESET_UNIQUE_HANDLE_KERNEL_LAUNCH,
        std::bind(&OpRemoteLookUpExecutor::TaskResetUniqueHandleKernelLaunch, this, std::placeholders::_1)},
    {OpTaskType::EVENT_WAIT_SEND_DONE,
        std::bind(&OpRemoteLookUpExecutor::TaskEventWaitSendDone, this, std::placeholders::_1)},
    {OpTaskType::RECV_VALUES_KERNEL_LAUNCH,
        std::bind(&OpRemoteLookUpExecutor::TaskRecvValuesKernelLaunch, this, std::placeholders::_1)},
    {OpTaskType::RECOVER_VALUE_KERNEL_LAUNCH,
        std::bind(&OpRemoteLookUpExecutor::TaskRecoverValueKernelLaunch, this, std::placeholders::_1)},
    {OpTaskType::REMOTE_LOOKUP_PROCESS,
        std::bind(&OpRemoteLookUpExecutor::TaskRemoteLookupProcess, this, std::placeholders::_1)},
    {OpTaskType::ISEND, std::bind(&OpRemoteLookUpExecutor::TaskIsend, this, std::placeholders::_1)},
    {OpTaskType::WAIT, std::bind(&OpRemoteLookUpExecutor::TaskWait, this, std::placeholders::_1)},
    {OpTaskType::WAIT_SEND_KEY_FINISH,
        std::bind(&OpRemoteLookUpExecutor::WaitSendKeyFinish, this, std::placeholders::_1)}
    }), mrManager_(mrManager), dispatcher_(dispatcher), localWorldRank_(localWorldRank)
{
}

OpRemoteLookUpExecutor::~OpRemoteLookUpExecutor()
{
    CHK_PRT(DeInit());
}

HcclResult OpRemoteLookUpExecutor::Init(std::map<u32, std::unique_ptr<TransportHeterog>> &transportHandleMap,
    std::map<u32, u32> &psRankIdMap)
{
    for (auto &iter : transportHandleMap) {
        if (iter.second != nullptr) {
            rankTransportMap_[iter.first] = iter.second.get();
        }
    }
    psRankId_.resize(psRankIdMap.size());
    psSize_ = psRankId_.size();
    for (auto &iter : psRankIdMap) {
        CHK_PRT_RET(iter.first == INVALID_UINT || iter.second == INVALID_UINT,
            HCCL_ERROR("[OpRemoteLookUpExecutor][Init]" \
            " executor init failed, ps[%u] rankId[%u], pls check rankTable's port and roleTable's port",
            iter.first, iter.second), HCCL_E_PARA);
        psRankId_[iter.first] = iter.second;
    }
    psIdOffset_.resize(psSize_);
    psValueStartAddr_.resize(psSize_);
    auto maxValue = *max_element(psRankId_.begin(), psRankId_.end());
    rankId2PsId_.insert(rankId2PsId_.begin(), maxValue + 1, 0);
    for (u32 i = 0; i < psSize_; i++) {
        rankId2PsId_[psRankId_[i]] = i;
    }
    return HCCL_SUCCESS;
}

HcclResult OpRemoteLookUpExecutor::DeInit()
{
    rankTransportMap_.clear();
    psRankId_.clear();
    if (defaultSendMem_.ptr() != nullptr) {
        if (!GetExternalInputHcclIsTcpMode()) {
            CHK_PRT(mrManager_->ReleaseKey(defaultSendMem_.ptr(), defaultSendMem_.size()));
        }
        defaultSendMem_.free();
    }

    if (defaultRecvMem_.ptr() != nullptr) {
        if (!GetExternalInputHcclIsTcpMode()) {
            CHK_PRT(mrManager_->ReleaseKey(defaultRecvMem_.ptr(), defaultRecvMem_.size()));
        }
        defaultRecvMem_.free();
    }

    if (psIdsMem_ != nullptr) {
        delete[] psIdsMem_;
        psIdsMem_ = nullptr;
    }

    return HCCL_SUCCESS;
}

HcclResult OpRemoteLookUpExecutor::ExecuteOp(HcclOpPtr opPtr)
{
    CHK_PTR_NULL(opPtr);
    HcclOperator<EmbeddingServiceParam> *ptr = static_cast<HcclOperator<EmbeddingServiceParam> *>(opPtr);
    if (ptr->param_.opType == EmbeddingServiceParam::OperatorType::ABORT_SELF) {
        blockFlag_ = false;
    } else if (ptr->param_.opType == EmbeddingServiceParam::OperatorType::REMOTE_LOOK_UP) {
        if (ptr->param_.keyMaxNum > keyMaxNum_) {
            if (defaultSendMem_.ptr() != nullptr) {
                if (!GetExternalInputHcclIsTcpMode()) {
                    CHK_RET(mrManager_->ReleaseKey(defaultSendMem_.ptr(), defaultSendMem_.size()));
                    CHK_RET(mrManager_->ReleaseKey(defaultRecvMem_.ptr(), defaultRecvMem_.size()));
                }
                defaultSendMem_.free();
                defaultRecvMem_.free();
            }

            if (psIdsMem_ != nullptr) {
                delete[] psIdsMem_;
                psIdsMem_ = nullptr;
            }

            u64 keyCount = ptr->param_.keyMaxNum * rankTransportMap_.size();
            u64 valueCount = ptr->param_.keyMaxNum * ptr->param_.embeddingParam.valueItemSize;
            HCCL_INFO("remote lookup keycount %llu, valueCount %llu, key max num %u, ps size %u", keyCount, valueCount,
                ptr->param_.keyMaxNum, rankTransportMap_.size());
            defaultSendMem_ = DeviceMem::alloc(sizeof(u64) * keyCount);
            defaultRecvMem_ = DeviceMem::alloc(valueCount + sizeof(u32));  // 申请存放错误码地址
            CHK_PTR_NULL(defaultSendMem_.ptr());
            CHK_PTR_NULL(defaultRecvMem_.ptr());
            if (!GetExternalInputHcclIsTcpMode()) {
                u32 lkeyForKeys = 0;
                u32 lkeyForValue = 0;
                CHK_RET(mrManager_->GetKey(defaultSendMem_.ptr(), sizeof(u64) * keyCount, lkeyForKeys));
                CHK_RET(mrManager_->GetKey(defaultRecvMem_.ptr(), valueCount + sizeof(u32), lkeyForValue));
            }
            keyMaxNum_ = ptr->param_.keyMaxNum;
            psIdsMem_ = new (std::nothrow) u32[keyMaxNum_];
            CHK_PTR_NULL(psIdsMem_);
            errorMsgAddr_ = static_cast<u8 *>(defaultRecvMem_.ptr()) + valueCount;
            CHK_PTR_NULL(errorMsgAddr_);
            *(static_cast<u32 *>(errorMsgAddr_)) = 0;
        }
        CHK_RET(ExecuteTask(ptr));
    } else {
        CHK_RET(ExecuteTask(ptr));
    }
    return HCCL_SUCCESS;
}

HcclResult OpRemoteLookUpExecutor::ExecuteTask(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    auto &taskQue = GetTaskQue<EmbeddingServiceParam>(opPtr);
    while (!taskQue.empty()) {
        HCCL_INFO("OpRemoteLookUpExecutor execute task: task type %u", taskQue.front()->taskType);
        CHK_RET(taskExeFuncTable_[taskQue.front()->taskType](taskQue.front()));
        taskQue.pop();
    }
    return HCCL_SUCCESS;
}

HcclResult OpRemoteLookUpExecutor::TaskMemcpy(OpTaskPtr &task)
{
    uint64_t beginTime = hrtMsprofSysCycleTime();
    void *dstMem = nullptr;
    u64 maxSize = 0;
    MemTask *taskPtr = dynamic_cast<MemTask *>(task.get());
    if (taskPtr->dstMem == nullptr) {
        dstMem = static_cast<u8 *>(defaultSendMem_.ptr()) + taskPtr->offset;
        maxSize = keyMaxNum_ * sizeof(u64);
    } else {
        dstMem = taskPtr->dstMem;
        maxSize = taskPtr->maxSize;
    }
    if (taskPtr->stream == nullptr) {
        HCCL_DEBUG("TaskMemcpy: mem cpy size %llu", taskPtr->size);
        TIME_PRINT(CHK_RET(hrtDrvMemCpy(dstMem, maxSize, taskPtr->srcMem, taskPtr->size)));
    } else {
        TIME_PRINT(CHK_RET(hrtMemAsyncCopy(dstMem, maxSize, taskPtr->srcMem, taskPtr->size,
            HcclRtMemcpyKind::HCCL_RT_MEMCPY_KIND_DEVICE_TO_DEVICE, taskPtr->stream)));
    }

    auto &profilingManager = ProfilingManager::Instance();
    CHK_RET(profilingManager.CallEsMsprofReportTaskApi(false, beginTime, ProfTaskType::TASK_LOOKUP_RESPONSE_MEMCPY));

    return HCCL_SUCCESS;
}

HcclResult OpRemoteLookUpExecutor::TaskSend(OpTaskPtr &task)
{
    TransTask *taskPtr = dynamic_cast<TransTask *>(task.get());

    TransportHeterog *heterogPtr = static_cast<TransportHeterog *>(GetTransport(taskPtr));
    void *sendMem = taskPtr->sendMem == nullptr ? (static_cast<u8 *>(defaultSendMem_.ptr()) + taskPtr->offset)
        : taskPtr->sendMem;
    taskPtr->sendMem = sendMem;

    CHK_RET(Send(taskPtr, heterogPtr));
    return HCCL_SUCCESS;
}

HcclResult OpRemoteLookUpExecutor::TaskRecv(OpTaskPtr &task)
{
    TransTask *taskPtr = dynamic_cast<TransTask *>(task.get());
    if (taskPtr->improbeRanks.size() > 0) {
        CHK_RET(TaskImprobeRecv(task));
        return HCCL_SUCCESS;
    }
    TransportHeterog *heterogPtr = static_cast<TransportHeterog *>(GetTransport(taskPtr));
    CHK_PTR_NULL(heterogPtr);
    TransportEndPointInfo srcEp(0, taskPtr->dstRank, taskPtr->tag);
    TransportEndPointInfo dstEp(0, taskPtr->srcRank, taskPtr->tag);
    TransportEndPointParam epParam(srcEp, dstEp);
    HcclMessageInfo *msg = nullptr;

    s32 flag = HCCL_IMPROBE_INCOMPLETED;
    HcclStatus status;

    uint64_t beginTime = hrtMsprofSysCycleTime();
    while (flag != HCCL_IMPROBE_COMPLETED) {
        if (!blockFlag_) {
            HCCL_INFO("TaskRecv execute abort self");
            return HCCL_E_AGAIN;
        }
        CHK_RET(heterogPtr->Improbe(epParam, flag, msg, status));
        if (flag != HCCL_IMPROBE_COMPLETED) {
            SaluSleep(LOOP_SLEEP_TIME_US);
        }
    }

    auto &profilingManager = hccl::ProfilingManager::Instance();
    CHK_RET(profilingManager.CallEsMsprofReportTaskApi(false, beginTime, ProfTaskType::TASK_EVENT_WAIT));

    u32 recvDataType = (msg->envelope.envelope.transData.dataType != HCCL_DATA_TYPE_RESERVED) ?
        msg->envelope.envelope.transData.dataType : msg->envelope.pcieEnvelope.dataType;
    TransportShmEvent *transportShm = dynamic_cast<TransportShmEvent *>(heterogPtr);
    u32 tableId = (transportShm != nullptr) ? msg->envelope.pcieEnvelope.tableId :\
        msg->envelope.envelope.transData.tableId;
    HcclUserRequire userRequire(tableId);
    void *recvMem = nullptr;
    if (msg->envelope.envelope.transData.errorFlag && status.count == sizeof(u32)) {
        recvMem = errorMsgAddr_;
    } else {
        recvMem = taskPtr->recvMem;
    }

    TransData recvData(reinterpret_cast<u64>(nullptr), reinterpret_cast<u64>(recvMem),
        status.count, static_cast<HcclDataType>(recvDataType));

    HCCL_DEBUG("Imrecv recvMem[%p], count[%d] dataType[%s]",
        reinterpret_cast<void *>(recvData.dstBuf), recvData.count, GetDataTypeEnumStr(recvData.dataType).c_str());
    HcclRequestInfo* request = nullptr;
    CHK_RET(heterogPtr->Imrecv(recvData, *msg, request));
    flag = HCCL_TEST_INCOMPLETED;
    HcclStatus compStatus;
    while (flag != HCCL_TEST_COMPLETED) {
        CHK_RET(heterogPtr->Test(*request, flag, compStatus));
        if (flag != HCCL_TEST_COMPLETED) {
            SaluSleep(LOOP_SLEEP_TIME_US);
        }
    }
    SetTaskOutputRegStatusInfo(taskPtr, status, recvDataType, userRequire);
    return HCCL_SUCCESS;
}

HcclResult OpRemoteLookUpExecutor::TaskImprobeRecv(OpTaskPtr &task)
{
    TransTask *taskPtr = dynamic_cast<TransTask *>(task.get());
    if (taskPtr->stream != nullptr) {
        return HCCL_E_PARA;
    }
    std::vector<TransportHeterog *> transports(taskPtr->improbeRanks.size(), nullptr);
    for (u32 i = 0; i < taskPtr->improbeRanks.size(); i++) {
        transports[i] = dynamic_cast<TransportHeterog *>(rankTransportMap_[taskPtr->improbeRanks[i]]);
        CHK_PTR_NULL(transports[i]);
    }

    TransportHeterog *heterogPtr = nullptr;
    HcclMessageInfo *msg = nullptr;
    HcclStatus status;
    TransportEndPointInfo srcEp(0, taskPtr->dstRank, taskPtr->tag);
    TransportEndPointInfo dstEp(0, taskPtr->srcRank, taskPtr->tag);
    TransportEndPointParam epParam(srcEp, dstEp);
    s32 flag = HCCL_IMPROBE_INCOMPLETED;

    uint64_t beginTime = hrtMsprofSysCycleTime();
    while (flag != HCCL_IMPROBE_COMPLETED && blockFlag_) {
        for (u32 i = 0; i < taskPtr->improbeRanks.size(); i++) {
            epParam.src.rank = taskPtr->improbeRanks[i];
            CHK_RET(transports[i]->Improbe(epParam, flag, msg, status));
            if (flag == HCCL_IMPROBE_COMPLETED) {
                heterogPtr = transports[i];
                tempTransQue_.push(heterogPtr);
                break;
            }
        }

        SaluSleep(LOOP_SLEEP_TIME_US);
    }

    HCCL_DEBUG("OpRemoteLookUpExecutor::TaskImprobe Recv start");
    if (heterogPtr == nullptr) {
        HCCL_INFO("TaskImprobeRecv execute abort self");
        return HCCL_E_AGAIN;
    }

    auto &profilingManager = hccl::ProfilingManager::Instance();
    CHK_RET(profilingManager.CallEsMsprofReportTaskApi(false, beginTime, ProfTaskType::TASK_EVENT_WAIT));
    HcclRequestInfo *request;
    u32 recvDataType = HCCL_DATA_TYPE_RESERVED;
    HcclUserRequire userRequire;
    GetProbeMsgInfo(heterogPtr, msg, recvDataType, userRequire);

    CHK_PRT_RET(recvDataType >= HCCL_DATA_TYPE_RESERVED, HCCL_ERROR("failed, recvDataType[%u] >= "
        "HCCL_DATA_TYPE_RESERVED[%d]", recvDataType, HCCL_DATA_TYPE_RESERVED), HCCL_E_PARA);

    TransData recvData(reinterpret_cast<u64>(nullptr), reinterpret_cast<u64>(taskPtr->recvMem),
        status.count, static_cast<HcclDataType>(recvDataType));

    HCCL_DEBUG("Imrecv srcRank[%d] dstRank[%d] recvMem[%llx], count[%d] dataType[%s] tableId[%u] globalStep[%lld]",
        epParam.src.rank, epParam.dst.rank, hash<u64>{}(recvData.dstBuf), recvData.count,
        GetDataTypeEnumStr(recvData.dataType).c_str(), userRequire.tableId, userRequire.globalStep);

    CHK_RET(heterogPtr->Imrecv(recvData, *msg, request));
    flag = HCCL_TEST_INCOMPLETED;
    HcclStatus compStatus;

    HCCL_DEBUG("OpRemoteLookUpExecutor::TaskImprobe Recv executed");
    while (flag != HCCL_TEST_COMPLETED && (recvData.count != 0)) {
        CHK_RET(heterogPtr->Test(*request, flag, compStatus));
        if (flag != HCCL_TEST_COMPLETED) {
            SaluSleep(LOOP_SLEEP_TIME_US);
        }
    }

    if (!taskPtr->outputs.empty()) {
        SetTaskOutputRegStatusInfo(taskPtr, status, recvDataType, userRequire);

        std::vector<void *> &outputs = taskPtr->outputs;
        if (outputs.size() > 1) {
            u32 dataTypeBytes = SIZE_TABLE[recvDataType];
            u64 *keyOffset = static_cast<u64 *>(outputs[OUTPUT_KEY_OFFSET_IDX]);
            *keyOffset = msg->envelope.pcieEnvelope.offset / dataTypeBytes;
            HCCL_INFO("TaskImprobeRecv success, keyOffset[%u]", *keyOffset);
            u64 *actualKeyCount = static_cast<u64 *>(outputs[OUTPUT_ACTUAL_KEY_COUNT_IDX]);
            // actualKeyCount为内部使用，等于实际key数量
            *actualKeyCount = status.count;
            *(static_cast<u32 *>(outputs[OUTPUT_WORKER_RANKID_IDX])) = epParam.src.rank;

            TransportHeterogRoce *transportRoce = dynamic_cast<TransportHeterogRoce *>(heterogPtr);
            if (transportRoce != nullptr) {
                RdmaBuffer *rdmaBuffer = reinterpret_cast<RdmaBuffer *>(msg->envelope.envelope.rsv);

                RdmaBuffer *rdmaBufferOut = reinterpret_cast<RdmaBuffer *>(outputs[OUTPUT_REPONSE_BUFFER_IDX]);

                CHK_PTR_NULL(rdmaBufferOut);
                *rdmaBufferOut = *rdmaBuffer;

                u64 &valueAddr = rdmaBufferOut->addr;
                u64 &valueAddrKey = rdmaBufferOut->buffKey;

                HCCL_DEBUG("transport is rdma recved valueAddr[%#llx], valueAddrKey[%llu]", valueAddr, valueAddrKey);
            } else {
                HCCL_DEBUG("transport is not rdma");
            }
        }
    }

    HCCL_DEBUG("OpRemoteLookUpExecutor::TaskImprobe Recv And Wait success");

    return HCCL_SUCCESS;
}

HcclResult OpRemoteLookUpExecutor::KeysSort(DropDuplicatesParams &params,
    std::map<u32, u32> &rankIdKeyNumMap)
{
    if (keyPosRecord_.size() < params.task.keyMaxNum) {
        keyPosRecord_.resize(params.task.keyMaxNum);
    }

    // 初始化psIdOffset_
    for (u32 psId = 0; psId < psSize_; ++psId) {
        psIdOffset_[psId] = 0;
    }
    if (params.task.intZerocpyFlag == ZERO_COPY_USED) {
        CHK_RET(hrtGetBatchPsIds((reinterpret_cast<u64 *>(*(params.task.keys))), &psIdsMem_, params.task.keyMaxNum));
    } else {
        CHK_RET(hrtGetBatchPsIds((reinterpret_cast<u64 *> (params.task.keys)), &psIdsMem_, params.task.keyMaxNum));
    }
    for (u32 i = 0; i < params.task.keyMaxNum; i++) {
            u32 psId = *(psIdsMem_ + i);
            u32 rankId = params.selpPtr->psRankId_[psId];

            keyPosRecord_[i].psId = psId;
            keyPosRecord_[i].pos = rankIdKeyNumMap[rankId];

#ifdef USE_AICORE_GATHERV2_INFER
            // rankIdAndPos信息供gather算子使用
            RankIdAndPos *rankIdAndPos = reinterpret_cast<RankIdAndPos *>(mappingMatrixS64_.Get());
            rankIdAndPos[i].pos = keyPosRecord_[i].pos;
            rankIdAndPos[i].rankId = rankId;
#endif
            rankIdKeyNumMap[rankId]++;
            s64 *dstMem = reinterpret_cast<s64 *>(static_cast<u8 *>(params.selpPtr->defaultSendMem_.ptr()) +
                params.task.keyMaxNum * KEY_SIZE * psId + psIdOffset_[psId]);
            if (params.task.intZerocpyFlag == ZERO_COPY_USED) {
                *dstMem = *((reinterpret_cast<u64 *>(*(params.task.keys))) + i);
            } else {
                *dstMem = *(params.task.keys + i);
            }

            psIdOffset_[psId] += KEY_SIZE;
            rankId2PsId_[rankId] = psId;
    }

    return HCCL_SUCCESS;
}

HcclResult OpRemoteLookUpExecutor::SendKeys(DropDuplicatesParams &params, std::map<u32, u32> &rankIdKeyNumMap,
    std::map<u32, u64> &waitLookUpRanks)
{
    std::map<TransportHeterog *, HcclRequestInfo *> requestMap;
    if (params.task.intZerocpyFlag == ZERO_COPY_USED) {
        CHK_PTR_NULL(params.task.tableIdAddr);
        params.task.tableIdAddr = reinterpret_cast<void *>(*(reinterpret_cast<u64 *>(params.task.tableIdAddr)));
    }
    CHK_PTR_NULL(params.task.tableIdAddr);
    params.task.tableId = *(static_cast<u32 *>(params.task.tableIdAddr));
    HCCL_INFO("RemoteLookUp SendKeys tag[%d] tableId[%u]", params.task.tag, params.task.tableId);
    for (auto &pair : rankIdKeyNumMap) {
        u32 rankId = pair.first;
        waitLookUpRanks[rankId] += pair.second * params.task.valueItemSize;

        HcclRequestInfo* request = nullptr;
        void *sendMem = static_cast<u8 *>(params.selpPtr->defaultSendMem_.ptr()) +
            rankId2PsId_[rankId] * params.task.keyMaxNum * KEY_SIZE;

        TransData sendData(reinterpret_cast<u64>(sendMem), reinterpret_cast<u64>(nullptr),
            pair.second * SIZE_TABLE[HCCL_DATA_TYPE_UINT64], HCCL_DATA_TYPE_INT8, false, params.task.tableId);

        TransportEndPointInfo srcEp(0, params.task.userRank, params.task.tag);
        TransportEndPointInfo dstEp(0, rankId, params.task.tag);
        TransportEndPointParam epParam(srcEp, dstEp);
        TransportHeterog *transPtr = dynamic_cast<TransportHeterog *>(params.selpPtr->rankTransportMap_[rankId]);
        uint64_t beginTime = hrtMsprofSysCycleTime();
        CHK_RET(transPtr->Isend(sendData, epParam, request));
        CHK_RET(ProfReportTransportShmEvent(transPtr, beginTime));
        requestMap[transPtr] = request;
    }

    bool allComplete = false;
    while (!allComplete) {
        for (auto &pair : requestMap) {
            if (pair.second == nullptr) {
                allComplete = true;
                continue;
            }

            s32 flag = HCCL_TEST_INCOMPLETED;
            HcclStatus compState;
            CHK_RET(pair.first->Test(*pair.second, flag, compState));
            if (flag == HCCL_TEST_COMPLETED) {
                pair.second = nullptr;
                allComplete = true;
            } else {
                allComplete = false;
            }
        }
        if (!allComplete) {
            SaluSleep(LOOP_SLEEP_TIME_US);
        }
    }
    return HCCL_SUCCESS;
}

HcclResult OpRemoteLookUpExecutor::ProfReportTransportShmEvent(TransportHeterog *transPtr, uint64_t beginTime) const
{
    if (dynamic_cast<TransportShmEvent*>(transPtr) != nullptr) {
        auto &profilingManager = ProfilingManager::Instance();
        CHK_RET(profilingManager.CallEsMsprofReportTaskApi(false, beginTime,
            ProfTaskType::TASK_SHARE_MEMORY_ISEND_RECORD));
        return HCCL_SUCCESS;
    }
    return HCCL_SUCCESS;
}

HcclResult OpRemoteLookUpExecutor::RecvValue(DropDuplicatesParams &params, std::map<u32, u64> &waitLookUpRanks)
{
    u64 offset = 0;
    CHK_PTR_NULL(params.selpPtr->defaultRecvMem_.ptr());
    for (auto &pair : waitLookUpRanks) {
        TransTask *transTaskPtr = new (std::nothrow) TransTask();
        CHK_PTR_NULL(transTaskPtr);

        transTaskPtr->count = pair.second / SIZE_TABLE[HCCL_DATA_TYPE_FP32];
        transTaskPtr->dataType = HCCL_DATA_TYPE_FP32; // embedding表数据为fp32
        transTaskPtr->recvMem = static_cast<u8 *>(params.selpPtr->defaultRecvMem_.ptr()) + offset;
        transTaskPtr->dstRank = pair.first;
        transTaskPtr->taskType = OpTaskType::RECV;
        transTaskPtr->outZerocpyFlag = params.task.outZerocpyFlag;
        transTaskPtr->tag = params.task.tag;
        psValueStartAddr_[rankId2PsId_[pair.first]] = transTaskPtr->recvMem;
        OpTaskPtr taskPtr(transTaskPtr);
        CHK_RET(params.selpPtr->TaskRecv(taskPtr));
#ifdef USE_AICORE_GATHERV2_INFER
        CHK_PRT_RET(params.task.valueItemSize == 0, HCCL_ERROR("valueItemSize is 0"), HCCL_E_PARA);
        rankIdAccumOffset_[pair.first] = offset / params.task.valueItemSize;
#endif
        offset += pair.second;
    }
    return HCCL_SUCCESS;
}

HcclResult OpRemoteLookUpExecutor::PrepareMappingMatrixForGatherValue(DropDuplicatesParams &params)
{
    RankIdAndPos *rankIdAndPos = reinterpret_cast<RankIdAndPos *>(params.selpPtr->mappingMatrixS64_.Get());

    for (u32 i = 0; i < params.selpPtr->mappingMatrixS64_.GetCount(); i++) {
        u32 rankId = rankIdAndPos[i].rankId;
        u32 pos = rankIdAndPos[i].pos;

        auto iter = params.selpPtr->rankIdAccumOffset_.find(rankId);
        CHK_PRT_RET(iter == params.selpPtr->rankIdAccumOffset_.end(),
            HCCL_ERROR("cannot find rankId[%u] in rankIdAccumOffset", rankId), HCCL_E_NOT_FOUND);

        params.selpPtr->mappingMatrixS64_.Get()[i] = static_cast<s64>(iter->second + pos);
    }

    return HCCL_SUCCESS;
}

HcclResult OpRemoteLookUpExecutor::RecoverValue(DropDuplicatesParams &params) const
{
    u32 valueCount = params.task.valueItemSize / SIZE_TABLE[HCCL_DATA_TYPE_FP32];
    void *outputMem = params.task.outZerocpyFlag == ZERO_COPY_USED ?
        reinterpret_cast<void *>(*(static_cast<u64 *>(params.task.value))) : params.task.value;
    for (u32 i = 0; i < params.task.keyMaxNum; i++) {
        void *srcMem = static_cast<u8 *>(psValueStartAddr_[keyPosRecord_[i].psId]) +
            keyPosRecord_[i].pos * params.task.valueItemSize;
        void *destMem = reinterpret_cast<u8*>(outputMem) + i * params.task.valueItemSize;
        FloatHighSpeedMove(static_cast<float *>(destMem), static_cast<float *>(srcMem), valueCount);
    }

    return HCCL_SUCCESS;
}
// 推理的流程
void OpRemoteLookUpExecutor::DropDuplicatesCallback(void *fnData)
{
    std::queue<DropDuplicatesParams> &que = *static_cast<std::queue<DropDuplicatesParams> *>(fnData);
    DropDuplicatesParams params = que.front();
    que.pop();
    que.push(params);

    std::map<u32, u32> rankIdKeyNumMap;
    // 按ps聚合，实际不排序
    // 当前是在ctlCpu上执行
    TIME_PRINT(CHK_RET_NULL(params.selpPtr->KeysSort(params, rankIdKeyNumMap)));

    // 只能用map
    std::map<u32, u64> waitLookUpRanks;
    TIME_PRINT(CHK_RET_NULL(params.selpPtr->SendKeys(params, rankIdKeyNumMap, waitLookUpRanks)));

    TIME_PRINT(CHK_RET_NULL(params.selpPtr->RecvValue(params, waitLookUpRanks)));

#ifdef USE_AICORE_GATHERV2_INFER
    // 为aicore算子准备映射矩阵信息
    TIME_PRINT(CHK_RET_NULL(params.selpPtr->PrepareMappingMatrixForGatherValue(params)));
    PRINT_ARRAY(params.selpPtr->mappingMatrixS64_.Get(), params.selpPtr->mappingMatrixS64_.GetCount(),
        "MappingMatrixS64");
#else
    TIME_PRINT(CHK_RET_NULL(params.selpPtr->RecoverValue(params)));
#endif

    PRINT_ARRAY((reinterpret_cast<s64 *>(*(params.task.keys))), params.task.keyMaxNum, "InferKeys");
    PRINT_ARRAY(static_cast<float *>(params.selpPtr->defaultRecvMem_.ptr()),
        params.task.keyMaxNum * (params.task.valueItemSize / SIZE_TABLE[HCCL_DATA_TYPE_FP32]), "InferRecvValue");

    HCCL_INFO("DropDuplicates success");
}

void OpRemoteLookUpExecutor::PrintValue(void *fnData)
{
    std::queue<DropDuplicatesParams> &que = *static_cast<std::queue<DropDuplicatesParams> *>(fnData);
    DropDuplicatesParams params = que.front();
    que.pop();
    que.push(params);

    PRINT_ARRAY(reinterpret_cast<float *>(params.task.value),
        params.task.keyMaxNum * (params.task.valueItemSize / SIZE_TABLE[HCCL_DATA_TYPE_FP32]), "InferRecoverValue");

    HCCL_INFO("PrintValue success");
}

HcclResult OpRemoteLookUpExecutor::TaskDropDuplicates(OpTaskPtr &task)
{
    PartitionMapTask *taskPtr = dynamic_cast<PartitionMapTask *>(task.get());
    if (taskPtr->keyMaxNum > keyMaxNum_) {
        if (defaultSendMem_.ptr() != nullptr) {
            if (!GetExternalInputHcclIsTcpMode()) {
                CHK_RET(mrManager_->ReleaseKey(defaultSendMem_.ptr(), defaultSendMem_.size()));
            }
            defaultSendMem_.free();
        }

        if (psIdsMem_ != nullptr) {
            delete[] psIdsMem_;
            psIdsMem_ = nullptr;
        }

        u64 keyCount = taskPtr->keyMaxNum * rankTransportMap_.size();
        defaultSendMem_ = DeviceMem::alloc(sizeof(u64) * keyCount);
        CHK_PTR_NULL(defaultSendMem_.ptr());
        if (!GetExternalInputHcclIsTcpMode()) {
            u32 lkey = 0;
            CHK_RET(mrManager_->GetKey(defaultSendMem_.ptr(), sizeof(u64) * keyCount, lkey));
        }
        keyMaxNum_ = taskPtr->keyMaxNum;
        psIdsMem_ = new (std::nothrow) u32[keyMaxNum_];
        CHK_PTR_NULL(psIdsMem_);
    }

    mappingMatrixS64_.Resize(taskPtr->keyMaxNum);

    DropDuplicatesParams params;
    params.task = *taskPtr;
    params.selpPtr = this;
    dropDuplicatesQue_.push(params);

    CHK_RET(hrtCallbackLaunch(&OpRemoteLookUpExecutor::DropDuplicatesCallback, &dropDuplicatesQue_,
        taskPtr->stream, true));

#ifdef USE_AICORE_GATHERV2_INFER
    CHK_RET(InferRecoverValueAicoreKernelLaunch(*taskPtr));
#ifdef DUMP_DATA
    CHK_RET(hrtCallbackLaunch(&OpRemoteLookUpExecutor::PrintValue, &dropDuplicatesQue_, taskPtr->stream, true));
#endif
#endif

    return HCCL_SUCCESS;
}

HcclResult OpRemoteLookUpExecutor::ChkServiceCancelKernelLaunch(const PartitionMapTask &mapTask)
{
    HdcsEmbeddingServiceParam para{};
    para.errorMsg = errorMsgAddr_;
    CHK_SAFETY_FUNC_RET(strcpy_s(para.groupName, GROUP_NAME_MAX_LEN, HCCL_WORLD_GROUP));

    rtKernelLaunchNames_t launchNames;
    launchNames.kernelName = "HdcsChkServiceCancel";
    launchNames.soName = "libccl_kernel.so";
    launchNames.opName = "ChkServiceCancel";
    CHK_RET(hrtAicpuKernelLaunch(launchNames, 1, &para, sizeof(HdcsEmbeddingServiceParam), nullptr,
        mapTask.stream));

    return HCCL_SUCCESS;
}

HcclResult OpRemoteLookUpExecutor::InferRecoverValueAicoreKernelLaunch(const PartitionMapTask &mapTask)
{
    s64 *mappingMatrix = this->mappingMatrixS64_.Get();
    u32 valueDim = mapTask.valueItemSize / SIZE_TABLE[HCCL_DATA_TYPE_FP32];

    // 避免宏函数为空时出现unused variable
    static_cast<void>(mappingMatrix);
    static_cast<void>(valueDim);

// 当前在HCCD场景下不能使用
#ifndef HCCD
    // 可重复调用，内部会避免重复初始化
    CHK_RET(aiCoreGatherV2_.Init());

    if (mapTask.outZerocpyFlag == ZERO_COPY_USED) {
        HCCL_ERROR("AiCore GatherV2 not supported!");
        return HCCL_E_NOT_SUPPORT;
    } else {
        void *outValue = mapTask.value;
        HcclDataType indicesType = HCCL_DATA_TYPE_INT64;

        CHK_RET(aiCoreGatherV2_.Run(this->defaultRecvMem_.ptr(), mappingMatrix, mapTask.keyMaxNum, mapTask.keyMaxNum,
            valueDim, indicesType, mapTask.stream, outValue));

        return HCCL_SUCCESS;
    }
#else
    HCCL_ERROR("AiCore GatherV2 not supported!");

    return HCCL_E_NOT_SUPPORT;
#endif
}

HcclResult OpRemoteLookUpExecutor::PrepareHostEnvelopeTemplates(OpTaskPtr &task)
{
    PartitionMapTask *taskPtr = dynamic_cast<PartitionMapTask *>(task.get());

    if (rdmaEnvelopeInfoHostMem_ == nullptr) {
        EXECEPTION_CATCH((rdmaEnvelopeInfoHostMem_ = make_unique<HcclEsRdmaInfoForLookup[]>(ES_MAX_PS_NUM)),
            return HCCL_E_PTR);
    }

    unique_ptr<HcclEsRdmaInfoForLookup[]> &rdmaEnveInfos = rdmaEnvelopeInfoHostMem_;

    taskPtr->haveRdmaConn = false;
    for (u32 i = 0; i < psSize_; ++i) {
        u32 rankId = psRankId_[i];
        u32 psId = i;

        TransportHeterog *&transportHeterog = rankTransportMap_[rankId];
        TransportShmEvent *transport = dynamic_cast<TransportShmEvent *>(transportHeterog);
        if (transport == nullptr) {
            TransportHeterogRoce *transportRoce = dynamic_cast<TransportHeterogRoce *>(transportHeterog);
            CHK_PRT_RET(transportRoce == nullptr,
                HCCL_ERROR("rankId[%u] transport type is not TransportShmEvent or TransportHeterogRoce.",
                rankId), HCCL_E_PARA);

            taskPtr->haveRdmaConn = true;
            rdmaEnveInfos[psId].commType = HcclHeterogCommType::RDMA;
            CHK_RET(FillEnvelope(task, rankId, rdmaEnveInfos[psId].envelope));

            HCCL_INFO("psId[%u], rank[%u] transport type is TransportHeterogRoce", psId, rankId);
        } else {
            rdmaEnveInfos[psId].commType = HcclHeterogCommType::PCIE;
            HCCL_INFO("psId[%u], rank[%u] transport type is TransportShmEvent", psId, rankId);
        }
    }

    return HCCL_SUCCESS;
}

HcclResult OpRemoteLookUpExecutor::FillEnvelope(OpTaskPtr &task, u32 rankId, HcclEnvelope &envelope)
{
    PartitionMapTask *taskPtr = dynamic_cast<PartitionMapTask *>(task.get());

    TransportHeterogRoce *transportRoce =
        dynamic_cast<TransportHeterogRoce *>(rankTransportMap_[rankId]);
    void *workerspaceAddr = taskPtr->keyTransferMem;

    u32 lkey = 0;
    CHK_RET(transportRoce->RegMr(workerspaceAddr, taskPtr->workerspaceMemSize, lkey, true));
    CHK_RET(transportRoce->RegMr(workerspaceAddr, taskPtr->workerspaceMemSize, lkey, false));

    HCCL_DEBUG("RegMr workerspace mem size[%llu], rdmaEnveInfosTransferMemSize[%llu]",
        taskPtr->workerspaceMemSize, taskPtr->rdmaEnveInfosTransferMemSize);

    TransportEndPointInfo srcEp(0, 0, taskPtr->tag);
    TransportEndPointInfo dstEp(0, rankId, taskPtr->tag);

    // envelope.transData 由device aicpu填写
    envelope.epParam.src = srcEp;
    envelope.epParam.dst = dstEp;
    envelope.protocol = 0;
    envelope.key = lkey;
    envelope.msn = 0;

    RdmaBuffer *rdmaBuffer = reinterpret_cast<RdmaBuffer *>(envelope.rsv);

    // addr内容 由device aicpu填写回写地址
    rdmaBuffer->addr = 0;
    // 填写回写地址的key
    rdmaBuffer->buffKey = lkey;

    return HCCL_SUCCESS;
}

HcclResult OpRemoteLookUpExecutor::AsyncCopyHostEnvelopeToDev(OpTaskPtr &task)
{
    PartitionMapTask *taskPtr = dynamic_cast<PartitionMapTask *>(task.get());

    if (!taskPtr->haveRdmaConn) {
        HCCL_INFO("haveRdmaConn is false, no need to AsyncCopyHostEnvelopeToDev");
        return HCCL_SUCCESS;
    }

    HcclEsRdmaInfoForLookup *devRdmaEnveInfos =
        static_cast<HcclEsRdmaInfoForLookup *>(taskPtr->rdmaEnveInfosTransferMem);
    unique_ptr<HcclEsRdmaInfoForLookup[]> &rdmaEnveInfos = rdmaEnvelopeInfoHostMem_;

    u64 maxMemSize = sizeof(HcclEsRdmaInfoForLookup) * ES_MAX_PS_NUM;
    if (devRdmaEnveInfos == nullptr || taskPtr->rdmaEnveInfosTransferMemSize < maxMemSize) {
        HCCL_ERROR("rdmaEnveInfosTransferMem is nullptr, or rdmaEnveInfosTransferMemSize[%llu] < maxMemSize[%llu]",
            taskPtr->rdmaEnveInfosTransferMemSize, maxMemSize);
        return HCCL_E_PARA;
    }

    u64 memSize = sizeof(HcclEsRdmaInfoForLookup) * psSize_;
    CHK_RET(hrtMemAsyncCopy(devRdmaEnveInfos, maxMemSize, rdmaEnveInfos.get(),
        memSize, HcclRtMemcpyKind::HCCL_RT_MEMCPY_KIND_HOST_TO_DEVICE, taskPtr->stream));

    PRINT_ARRAY(reinterpret_cast<u64 *>(rdmaEnveInfos.get()), memSize / sizeof(u64), "LookupHostRdmaEnveInfos");

    HCCL_DEBUG("AsyncCopyHostEnvelopeToDev SUCCESS !");
    return HCCL_SUCCESS;
}

HcclResult OpRemoteLookUpExecutor::TaskSendKeysByRdma(OpTaskPtr &task)
{
    PartitionMapTask *taskPtr = dynamic_cast<PartitionMapTask *>(task.get());
    HcclEsRdmaInfoForLookup *devRdmaEnveInfos =
        static_cast<HcclEsRdmaInfoForLookup *>(taskPtr->rdmaEnveInfosTransferMem);
    unique_ptr<HcclEsRdmaInfoForLookup[]> &rdmaEnveInfos = rdmaEnvelopeInfoHostMem_;

    for (u32 i = 0; i < psSize_; ++i) {
        u32 rankId = psRankId_[i];
        u32 psId = i;

        if (rdmaEnveInfos[psId].commType != HcclHeterogCommType::RDMA) {
            HCCL_DEBUG("rankId[%u], psId[%u] send transport type is not rdma, commType[%d]",
                rankId, psId, rdmaEnveInfos[psId].commType);
            continue;
        }

        HCCL_DEBUG("rankId[%u], psId[%u] send transport type is rdma", rankId, psId);

        HcclEnvelope &devEnvelope = devRdmaEnveInfos[psId].envelope;
        TransportHeterogRoce *transportRoce = dynamic_cast<TransportHeterogRoce *>(rankTransportMap_[rankId]);

        HcclResult ret = transportRoce->SendEnvelope(devEnvelope, taskPtr->stream);
        if (ret == HCCL_E_TIMEOUT) {
            ExceptionTransportInfoRetriever(localWorldRank_, rankTransportMap_);
        }
        CHK_RET(ret);
    }

    return HCCL_SUCCESS;
}

HcclResult OpRemoteLookUpExecutor::TaskRemoteLookupProcess(OpTaskPtr &task)
{
    CHK_RET(PrepareHostEnvelopeTemplates(task));

    // 去重
    CHK_RET(TaskKeyDropDuplicatesKernelLaunch(task));

    CHK_RET(AsyncCopyHostEnvelopeToDev(task));

    // send key
    CHK_RET(TaskSendKeysKernelLaunch(task));

    CHK_RET(TaskSendKeysByRdma(task));

    // notify处理
    CHK_RET(TaskSendKeysRecord(task));

    CHK_RET(TaskEventWaitRecvDone(task));

    CHK_RET(TaskResetUniqueHandleKernelLaunch(task));

    CHK_RET(TaskEventWaitSendDone(task));
    // recv value
    CHK_RET(TaskRecvValuesKernelLaunch(task));

    // recover value
#ifdef USE_AICORE_GATHERV2
    CHK_RET(RecoverValueAicoreKernelLaunch(task));
#ifdef TIME_PROFILING
    // gather结束后统计时间
    CHK_RET(TaskGatherFinishKernelLaunch(task));
#endif
#else
    CHK_RET(TaskRecoverValueKernelLaunch(task));
#endif
    return HCCL_SUCCESS;
}

EsLoopUpPara OpRemoteLookUpExecutor::CreateEsLoopUpPara(const PartitionMapTask* taskPtr) const
{
    EsLoopUpPara esLoopUpPara = {
        .tag = taskPtr->tag,
        .srcAddr = taskPtr->keys,
        .dstAddr = taskPtr->value,
        .dataSize = static_cast<u32>(taskPtr->keyMaxNum)
    };
    return esLoopUpPara;
}

// RecoverValue下发到Aicore的gatherV2算子进行执行
HcclResult OpRemoteLookUpExecutor::RecoverValueAicoreKernelLaunch(OpTaskPtr &task)
{
    uint64_t beginTime = hrtMsprofSysCycleTime();
    PartitionMapTask *taskPtr = dynamic_cast<PartitionMapTask *>(task.get());
    if (taskPtr->outZerocpyFlag == ZERO_COPY_USED) {
        HCCL_ERROR("AiCore GatherV2 not supported out zero copy!");
        return HCCL_E_NOT_SUPPORT;
    }

    vector<HcclRtNotify> &notifyInfo = taskPtr->notifyInfo;
    auto &profilingManager = hccl::ProfilingManager::Instance();
    EsLoopUpPara esLoopUpPara = CreateEsLoopUpPara(taskPtr);
    CHK_RET(profilingManager.CallMsprofReportAdditionInfoForEsLookup(esLoopUpPara,
            ProfTaskType::TASK_RECOVER_VALUE_AICORE));
    void *stream{};
    if (taskPtr->usePipeline) {
        // 等待从流1 recv values done
        s32 waitNotifyIndex = taskPtr->cubeIndex * ES_PIPELINE_NOFITY_COUNT + PIPELINE_RECOVER_VALUE_NOTIFY_INDEX - 1;
        stream = taskPtr->subStreamInfo[PIPELINE_RECOVER_VALUE_SUB_STREAM_INDEX];
        Stream tmpStream(stream);

        CHK_RET(static_cast<LocalNotify*>(notifyInfo[waitNotifyIndex])->Wait(tmpStream, dispatcher_));
    } else {
        stream = taskPtr->stream;
    }

// 当前在HCCD场景下不能使用
#ifndef HCCD
    // 可重复调用，内部会避免重复初始化
    CHK_RET(aiCoreGatherV2_.Init());

    void *mappingMatrix{};
    if (taskPtr->pairedMode) {
        mappingMatrix = taskPtr->indices;
    } else {
        mappingMatrix = static_cast<s64 *>(taskPtr->keyTransferMem) + taskPtr->keyMaxNum;
        if (taskPtr->enableKeyCounter) {
            mappingMatrix = static_cast<s8 *>(mappingMatrix) + taskPtr->keyMaxNum * ES_KEY_COUNTER_MEM_BYTES_SIZE;
        }
    }

    u32 valueDim = taskPtr->valueItemSize / SIZE_TABLE[HCCL_DATA_TYPE_FP32];
    HcclDataType indicesType = HCCL_DATA_TYPE_INT32;

    CHK_RET(aiCoreGatherV2_.Run(taskPtr->valueTransferMem, mappingMatrix, taskPtr->keyMaxNum, taskPtr->keyMaxNum,
        valueDim, indicesType, stream, taskPtr->value));

    CHK_RET(profilingManager.CallEsMsprofReportTaskApi(false, beginTime, ProfTaskType::TASK_RECOVER_VALUE_AICORE));
    return HCCL_SUCCESS;
#else
    HCCL_ERROR("AiCore GatherV2 not supported!");
    CHK_RET(profilingManager.CallEsMsprofReportTaskApi(false, beginTime, ProfTaskType::TASK_RECOVER_VALUE_AICORE));
    return HCCL_E_NOT_SUPPORT;
#endif
}

HcclResult OpRemoteLookUpExecutor::TaskKeyDropDuplicatesKernelLaunch(OpTaskPtr &task)
{
    uint64_t beginTime = hrtMsprofSysCycleTime();
    PartitionMapTask *taskPtr = dynamic_cast<PartitionMapTask *>(task.get());
    std::vector<HcclRtNotify> &notifyInfo = taskPtr->notifyInfo;
    HdcsEmbeddingServiceParam para{};
    SetupTaskHdcsEmbeddingServiceParam(taskPtr, para);

    rtKernelLaunchNames_t launchNames;
    launchNames.kernelName = "HdcsCollLookupKeysDuplicates";
    launchNames.soName = "libccl_kernel.so";
    launchNames.opName = "CollLookupKeysDuplicates";
    auto &profilingManager = hccl::ProfilingManager::Instance();
    EsLoopUpPara esLoopUpPara = CreateEsLoopUpPara(taskPtr);
    CHK_RET(profilingManager.CallMsprofReportAdditionInfoForEsLookup(esLoopUpPara,
            ProfTaskType::TASK_KEY_DROP_DUPLICATES));
    CHK_RET(hrtAicpuKernelLaunch(launchNames, 1, &para, sizeof(HdcsEmbeddingServiceParam), nullptr, taskPtr->stream));

    if (taskPtr->usePipeline) {
        // 通知从流0开始send key
        s32 recordNotifyIndex = taskPtr->cubeIndex * ES_PIPELINE_NOFITY_COUNT;
        Stream tmpStream(taskPtr->stream);
        CHK_RET(static_cast<LocalNotify*>(notifyInfo[recordNotifyIndex])->Post(tmpStream, dispatcher_));
    }
    CHK_RET(profilingManager.CallEsMsprofReportTaskApi(false, beginTime, ProfTaskType::TASK_KEY_DROP_DUPLICATES));
    return HCCL_SUCCESS;
}

void OpRemoteLookUpExecutor::SetupTaskHdcsEmbeddingServiceParam(PartitionMapTask *taskPtr,
    HdcsEmbeddingServiceParam &para) const
{
    para.tag = taskPtr->tag;
    para.keys = taskPtr->keys;
    para.keyMaxNum = taskPtr->keyMaxNum;
    para.keyType = HCCL_DATA_TYPE_INT64;
    para.values = taskPtr->value;
    para.valueCount = taskPtr->keyMaxNum * taskPtr->valueItemSize;
    para.valuesType = HCCL_DATA_TYPE_FP32;
    para.tableIdAddr =  taskPtr->outputTableId;
    para.tableId = taskPtr->tableId;
    para.insertFlag = taskPtr->insertFlag;
    para.valueItemSize = taskPtr->valueItemSize;
    para.keyTransferMem = taskPtr->keyTransferMem;
    para.keyTransferMemSize = taskPtr->keyTransferMemSize;
    para.valueTransferMem = taskPtr->valueTransferMem;
    para.valueTransferMemSize = taskPtr->valueTransferMemSize;
    para.rdmaEnveInfosTransferMem = taskPtr->rdmaEnveInfosTransferMem;
    para.rdmaEnveInfosTransferMemSize = taskPtr->rdmaEnveInfosTransferMemSize;

    para.cubeIndex = taskPtr->cubeIndex;
    para.pipelineKeyNum = taskPtr->pipelineKeyNum;
    para.pairedMode = taskPtr->pairedMode;
    para.usePipeline = taskPtr->usePipeline;
    para.enableKeyCounter = taskPtr->enableKeyCounter;

    para.intZerocpyFlag = taskPtr->intZerocpyFlag;
    para.outZerocpyFlag = taskPtr->outZerocpyFlag;
    para.disableUnique = taskPtr->disableUnique;

    para.uniqued = taskPtr->uniqued;
    para.keyNumInput = taskPtr->keyNumInput;
    para.uniqueIndices = taskPtr->uniqueIndices;
    para.keyCount = taskPtr->keyCount;
    para.haveRdmaConn = taskPtr->haveRdmaConn;

    if (para.pairedMode) {
        para.indices = taskPtr->indices;
        para.numUniqued = taskPtr->numUniqued;
        para.psSeg = taskPtr->psSeg;
        para.psSegNum = taskPtr->psSegNum;
    }
}

HcclResult OpRemoteLookUpExecutor::TaskSendKeysKernelLaunch(OpTaskPtr &task)
{
    uint64_t beginTime = hrtMsprofSysCycleTime();
    PartitionMapTask *taskPtr = dynamic_cast<PartitionMapTask *>(task.get());
    std::vector<rtStream_t> &subStreamInfo = taskPtr->subStreamInfo;
    std::vector<HcclRtNotify> &notifyInfo = taskPtr->notifyInfo;

    if (taskPtr->usePipeline) {
        // 等待主流去重done
        s32 waitNotifyIndex = taskPtr->cubeIndex * ES_PIPELINE_NOFITY_COUNT + PIPELINE_SEND_KEYS_NOTIFY_INDEX - 1;
        Stream tmpStream(subStreamInfo[PIPELINE_SEND_KEYS_SUB_STREAM_INDEX]);
        CHK_RET(static_cast<LocalNotify*>(notifyInfo[waitNotifyIndex])->Wait(tmpStream, dispatcher_));
    }

    HdcsLookupPipelineParams para{};
    para.tag = taskPtr->tag;
    para.cubeIndex = taskPtr->cubeIndex;
    para.usePipeline = taskPtr->usePipeline;
    para.tableIdAddr = taskPtr->outputTableId;
    para.intZerocpyFlag = taskPtr->intZerocpyFlag;

    rtKernelLaunchNames_t launchNames;
    launchNames.kernelName = "HdcsCollLookupSendKeys";
    launchNames.soName = "libccl_kernel.so";
    launchNames.opName = "CollLookupSendKeys";

    auto &profilingManager = hccl::ProfilingManager::Instance();
    EsLoopUpPara esLoopUpPara = CreateEsLoopUpPara(taskPtr);
    CHK_RET(profilingManager.CallMsprofReportAdditionInfoForEsLookup(esLoopUpPara, ProfTaskType::TASK_SEND_KEYS));

    if (!taskPtr->usePipeline) {
        CHK_RET(hrtAicpuKernelLaunch(launchNames, 1, &para, sizeof(HdcsLookupPipelineParams), nullptr,
            taskPtr->stream));

        CHK_RET(profilingManager.CallEsMsprofReportTaskApi(false, beginTime, ProfTaskType::TASK_SEND_KEYS));
        return HCCL_SUCCESS;
    }

    CHK_RET(hrtAicpuKernelLaunch(launchNames, 1, &para, sizeof(HdcsLookupPipelineParams), nullptr,
        subStreamInfo[PIPELINE_SEND_KEYS_SUB_STREAM_INDEX]));

    // 通知从流1开始recv values
    s32 recordNotifyIndex = taskPtr->cubeIndex * ES_PIPELINE_NOFITY_COUNT + PIPELINE_SEND_KEYS_NOTIFY_INDEX;
    Stream tmpStream(subStreamInfo[PIPELINE_SEND_KEYS_SUB_STREAM_INDEX]);
    CHK_RET(static_cast<LocalNotify*>(notifyInfo[recordNotifyIndex])->Post(tmpStream, dispatcher_));

    CHK_RET(profilingManager.CallEsMsprofReportTaskApi(false, beginTime, ProfTaskType::TASK_SEND_KEYS));
    return HCCL_SUCCESS;
}

HcclResult OpRemoteLookUpExecutor::TaskSendKeysRecord(OpTaskPtr &task)
{
    uint64_t beginTime = hrtMsprofSysCycleTime();
    PartitionMapTask *taskPtr = dynamic_cast<PartitionMapTask *>(task.get());

    auto &profilingmanager = hccl::ProfilingManager::Instance();
    EsLoopUpPara esLoopUpPara = CreateEsLoopUpPara(taskPtr);
    CHK_RET(profilingmanager.CallMsprofReportAdditionInfoForEsLookup(esLoopUpPara,
            ProfTaskType::TASK_SEND_KEYS_RECORD));

    unique_ptr<HcclEsRdmaInfoForLookup[]> &rdmaEnveInfos = rdmaEnvelopeInfoHostMem_;
    if (!taskPtr->usePipeline) {
        std::shared_ptr<RemoteNotify> isendDoneNotify = nullptr;
        for (u32 i = 0; i < psSize_; ++i) {
            u32 rankId = psRankId_[i];
            // rdma不需要通知
            if (rdmaEnveInfos[i].commType != HcclHeterogCommType::PCIE) {
                HCCL_DEBUG("rankId[%u], psId[%u] transport type[%d] is not pcie, no need to sendkey record",
                    rankId, i, rdmaEnveInfos[i].commType);
                continue;
            }

            TransportShmEvent *transport = dynamic_cast<TransportShmEvent *>(rankTransportMap_[rankId]);
            if (transport == nullptr) {
                HCCL_ERROR("rankId[%u], psId[%u] transport type is not TransportShmEvent", rankId, i);
                return HCCL_E_INTERNAL;
            }
            CHK_RET(transport->GetIsendDoneSignal(isendDoneNotify));
            // isendDoneNotify
            Stream tmpStream(taskPtr->stream);
            isendDoneNotify->Post(tmpStream, dispatcher_);
        }
    }
    CHK_RET(profilingmanager.CallEsMsprofReportTaskApi(false, beginTime, ProfTaskType::TASK_SEND_KEYS_RECORD));
    return HCCL_SUCCESS;
}

u32 OpRemoteLookUpExecutor::GetTimeOut()
{
    HCCL_DEBUG("GetTimeOut start.");
    u32 defTimeOut = NOTIFY_DEFAULT_WAIT_TIME;
    if (GetExternalInputHcclExecTimeoutSet() != HcclExecTimeoutSet::HCCL_EXEC_TIMEOUT_NOT_SET) {
        HCCL_DEBUG("GetTimeOut from external input.");
        defTimeOut = GetExternalInputHcclExecTimeOut();
    }

    if (defTimeOut == 0) {
        HCCL_DEBUG("GetTimeOut get external timeout is 0.");
    }

    HCCL_DEBUG("timeout is %u.", defTimeOut);
    return defTimeOut;
}

HcclResult OpRemoteLookUpExecutor::TaskEventWaitRecvDone(OpTaskPtr &task)
{
    uint64_t beginTime = hrtMsprofSysCycleTime();
    PartitionMapTask *taskPtr = dynamic_cast<PartitionMapTask *>(task.get());

    auto &profilingmanager = hccl::ProfilingManager::Instance();
    EsLoopUpPara esLoopUpPara = CreateEsLoopUpPara(taskPtr);
    CHK_RET(profilingmanager.CallMsprofReportAdditionInfoForEsLookup(esLoopUpPara,
            ProfTaskType::TASK_EVENT_WAIT_RECV_DONE));

    unique_ptr<HcclEsRdmaInfoForLookup[]> &rdmaEnveInfos = rdmaEnvelopeInfoHostMem_;
    if (!taskPtr->usePipeline) {
        std::shared_ptr<LocalIpcNotify> remoteImrecvDoneNotify = nullptr;
        for (u32 i = 0; i < psSize_; ++i) {
            u32 rankId = psRankId_[i];

            // rdma场景：ps未做notify通知，worker也不等待。
            if (rdmaEnveInfos[i].commType != HcclHeterogCommType::PCIE) {
                HCCL_DEBUG("rankId[%u], psId[%u] transport type[%d] is not pcie, no need to WaitRecvDone",
                    rankId, i, rdmaEnveInfos[i].commType);
                continue;
            }

            TransportShmEvent *transport = dynamic_cast<TransportShmEvent *>(rankTransportMap_[rankId]);
            CHK_RET(transport->GetRemoteImrecvDoneSignal(remoteImrecvDoneNotify));

            // remoteImrecvDoneNotify
            Stream tmpStream(taskPtr->stream);
            remoteImrecvDoneNotify->Wait(tmpStream, dispatcher_, INVALID_VALUE_STAGE, GetTimeOut());
        }
    }

    CHK_RET(profilingmanager.CallEsMsprofReportTaskApi(false, beginTime, ProfTaskType::TASK_EVENT_WAIT_RECV_DONE));

    return HCCL_SUCCESS;
}

HcclResult OpRemoteLookUpExecutor::TaskEventWaitSendDone(OpTaskPtr &task)
{
    uint64_t beginTime = hrtMsprofSysCycleTime();
    PartitionMapTask *taskPtr = dynamic_cast<PartitionMapTask *>(task.get());
    auto &profilingManager = hccl::ProfilingManager::Instance();
    EsLoopUpPara esLoopUpPara = CreateEsLoopUpPara(taskPtr);
    CHK_RET(profilingManager.CallMsprofReportAdditionInfoForEsLookup(esLoopUpPara,
            ProfTaskType::TASK_EVENT_WAIT_SEND_DONE));
    if (!taskPtr->usePipeline) {
        std::shared_ptr<LocalIpcNotify> remoteIsendDoneNotify = nullptr;
        for (u32 i = 0; i < psSize_; ++i) {
            u32 rankId = psRankId_[i];
            TransportHeterog *transport = rankTransportMap_[rankId];
            CHK_RET(transport->GetRemoteIsendDoneSignal(remoteIsendDoneNotify));
            Stream tmpStream(taskPtr->stream);
            TransportHeterogRoce *tmpTransport = dynamic_cast<TransportHeterogRoce*>(transport);
            if (tmpTransport != nullptr && tmpTransport->tagQpInfo_.qpMode == OPBASE_QP_MODE_EXT) {
                auto meta = HcclOpMetaInfo::GetOneForRecieve();
                CHK_RET(InitTask(dispatcher_, tmpStream, meta.isEnableCache, meta.GetCacheKey()));
                remoteIsendDoneNotify->Wait(tmpStream, dispatcher_, INVALID_VALUE_STAGE, GetTimeOut());
                CHK_RET(LaunchTask(dispatcher_, tmpStream));
            } else {
                remoteIsendDoneNotify->Wait(tmpStream, dispatcher_, INVALID_VALUE_STAGE, GetTimeOut());
            }
        }
    }

    CHK_RET(profilingManager.CallEsMsprofReportTaskApi(false, beginTime, ProfTaskType::TASK_EVENT_WAIT_SEND_DONE));
    return HCCL_SUCCESS;
}

HcclResult OpRemoteLookUpExecutor::TaskRecvValuesKernelLaunch(OpTaskPtr &task)
{
    uint64_t beginTime = hrtMsprofSysCycleTime();
    PartitionMapTask *taskPtr = dynamic_cast<PartitionMapTask *>(task.get());
    std::vector<rtStream_t> &subStreamInfo = taskPtr->subStreamInfo;
    std::vector<HcclRtNotify> &notifyInfo = taskPtr->notifyInfo;
    auto &profilingManager = hccl::ProfilingManager::Instance();
    EsLoopUpPara esLoopUpPara = CreateEsLoopUpPara(taskPtr);
    CHK_RET(profilingManager.CallMsprofReportAdditionInfoForEsLookup(esLoopUpPara, ProfTaskType::TASK_RECV_VALUES));
    if (taskPtr->usePipeline) {
        // 等待从流0 send keys done
        s32 waitNotifyIndex = taskPtr->cubeIndex * ES_PIPELINE_NOFITY_COUNT + PIPELINE_RECV_VALUES_NOTIFY_INDEX - 1;
        Stream tmpStream(subStreamInfo[PIPELINE_RECV_VALUES_SUB_STREAM_INDEX]);
        CHK_RET(static_cast<LocalNotify*>(notifyInfo[waitNotifyIndex])->Wait(tmpStream, dispatcher_));
    }

    HdcsLookupPipelineParams para{};
    para.tag = taskPtr->tag;
    para.cubeIndex = taskPtr->cubeIndex;
    para.usePipeline = taskPtr->usePipeline;
    para.tableIdAddr = taskPtr->outputTableId;
    para.intZerocpyFlag = taskPtr->intZerocpyFlag;

    rtKernelLaunchNames_t launchNames;
    launchNames.kernelName = "HdcsCollLookupRecvValues";
    launchNames.soName = "libccl_kernel.so";
    launchNames.opName = "CollLookupRecvValues";

    if (!taskPtr->usePipeline) {
        CHK_RET(hrtAicpuKernelLaunch(launchNames, 1, &para, sizeof(HdcsLookupPipelineParams), nullptr,
            taskPtr->stream));

        CHK_RET(profilingManager.CallEsMsprofReportTaskApi(false, beginTime, ProfTaskType::TASK_RECV_VALUES));
        return HCCL_SUCCESS;
    }

    CHK_RET(hrtAicpuKernelLaunch(launchNames, 1, &para, sizeof(HdcsLookupPipelineParams), nullptr,
        subStreamInfo[PIPELINE_RECV_VALUES_SUB_STREAM_INDEX]));

    // 通知从流2开始recover value
    s32 recordNotifyIndex = taskPtr->cubeIndex * ES_PIPELINE_NOFITY_COUNT + PIPELINE_RECV_VALUES_NOTIFY_INDEX;
    Stream tmpStream(subStreamInfo[PIPELINE_RECV_VALUES_SUB_STREAM_INDEX]);
    CHK_RET(static_cast<LocalNotify*>(notifyInfo[recordNotifyIndex])->Post(tmpStream, dispatcher_));

    CHK_RET(profilingManager.CallEsMsprofReportTaskApi(false, beginTime, ProfTaskType::TASK_RECV_VALUES));
    return HCCL_SUCCESS;
}

HcclResult OpRemoteLookUpExecutor::TaskRecoverValueKernelLaunch(OpTaskPtr &task)
{
    PartitionMapTask *taskPtr = dynamic_cast<PartitionMapTask *>(task.get());
    std::vector<rtStream_t> &subStreamInfo = taskPtr->subStreamInfo;
    std::vector<HcclRtNotify> &notifyInfo = taskPtr->notifyInfo;

    if (taskPtr->usePipeline) {
        // 等待从流1 recv values done
        s32 waitNotifyIndex = taskPtr->cubeIndex * ES_PIPELINE_NOFITY_COUNT + PIPELINE_RECOVER_VALUE_NOTIFY_INDEX - 1;
        Stream tmpStream(subStreamInfo[PIPELINE_RECOVER_VALUE_SUB_STREAM_INDEX]);
        CHK_RET(static_cast<LocalNotify*>(notifyInfo[waitNotifyIndex])->Wait(tmpStream, dispatcher_));
    }

    HdcsLookupPipelineParams para{};
    para.tag = taskPtr->tag;
    para.cubeIndex = taskPtr->cubeIndex;
    para.usePipeline = taskPtr->usePipeline;
    para.tableIdAddr = taskPtr->outputTableId;
    para.intZerocpyFlag = taskPtr->intZerocpyFlag;

    rtKernelLaunchNames_t launchNames;
    launchNames.kernelName = "HdcsCollLookupRecoverValue";
    launchNames.soName = "libccl_kernel.so";
    launchNames.opName = "CollLookupRecoverValue";
    if (taskPtr->usePipeline) {
        CHK_RET(hrtAicpuKernelLaunch(launchNames, 1, &para, sizeof(HdcsLookupPipelineParams), nullptr,
            subStreamInfo[PIPELINE_RECOVER_VALUE_SUB_STREAM_INDEX]));
    } else {
        CHK_RET(hrtAicpuKernelLaunch(launchNames, 1, &para, sizeof(HdcsLookupPipelineParams), nullptr,
            taskPtr->stream));
    }
    return HCCL_SUCCESS;
}

HcclResult OpRemoteLookUpExecutor::TaskGatherFinishKernelLaunch(OpTaskPtr &task)
{
    uint64_t beginTime = hrtMsprofSysCycleTime();
    PartitionMapTask *taskPtr = dynamic_cast<PartitionMapTask *>(task.get());

    // 流水场景下需要调整stream
    HdcsLookupPipelineParams para{};
    para.tag = taskPtr->tag;
    para.tableIdAddr = taskPtr->outputTableId;
    para.intZerocpyFlag = taskPtr->intZerocpyFlag;

    rtKernelLaunchNames_t launchNames;
    launchNames.kernelName = "HdcsCollLookupGatherFinish";
    launchNames.soName = "libccl_kernel.so";
    launchNames.opName = "CollGatherFinish";
    auto &profilingManager = hccl::ProfilingManager::Instance();
    EsLoopUpPara esLoopUpPara = CreateEsLoopUpPara(taskPtr);
    CHK_RET(profilingManager.CallMsprofReportAdditionInfoForEsLookup(esLoopUpPara, ProfTaskType::TASK_GATHER_FINISH));
    CHK_RET(hrtAicpuKernelLaunch(launchNames, 1, &para, sizeof(HdcsLookupPipelineParams), nullptr,
        taskPtr->stream));
    CHK_RET(profilingManager.CallEsMsprofReportTaskApi(false, beginTime, ProfTaskType::TASK_GATHER_FINISH));
    return HCCL_SUCCESS;
}

HcclResult OpRemoteLookUpExecutor::TaskResetUniqueHandleKernelLaunch(OpTaskPtr &task)
{
    uint64_t beginTime = hrtMsprofSysCycleTime();
    PartitionMapTask *taskPtr = dynamic_cast<PartitionMapTask *>(task.get());
    std::vector<rtStream_t> &subStreamInfo = taskPtr->subStreamInfo;
    std::vector<HcclRtNotify> &notifyInfo = taskPtr->notifyInfo;

    HdcsLookupPipelineParams para{};
    para.tag = taskPtr->tag;
    para.cubeIndex = taskPtr->cubeIndex;
    para.usePipeline = taskPtr->usePipeline;
    para.tableIdAddr = taskPtr->outputTableId;
    para.intZerocpyFlag = taskPtr->intZerocpyFlag;

    rtKernelLaunchNames_t launchNames;
    launchNames.kernelName = "HdcsCollLookupResetUniqueHandle";
    launchNames.soName = "libccl_kernel.so";
    launchNames.opName = "CollLookupResetUniqueHandle";
    auto &profilingManager = hccl::ProfilingManager::Instance();
    EsLoopUpPara esLoopUpPara = CreateEsLoopUpPara(taskPtr);
    CHK_RET(profilingManager.CallMsprofReportAdditionInfoForEsLookup(esLoopUpPara,
            ProfTaskType::TASK_RESET_UNIQUE_HANDLE));
    CHK_RET(hrtAicpuKernelLaunch(launchNames, 1, &para, sizeof(HdcsLookupPipelineParams), nullptr, taskPtr->stream));

    if (taskPtr->usePipeline) {
        CHK_PRT_RET(notifyInfo.size() < NOTIFY_RESET_UNIQUE_REVERSE_IDX, HCCL_ERROR("notifyInfo.size < %u",
            NOTIFY_RESET_UNIQUE_REVERSE_IDX), HCCL_E_PARA);
        // 等待 lookup流水 done，notify array中倒数第二个notify用于HdcsCollLookupResetUniqueHandle算子执行后的同步
        s32 waitNotifyIndex = notifyInfo.size() - NOTIFY_RESET_UNIQUE_REVERSE_IDX;
        Stream waitStream(taskPtr->stream);
        CHK_RET(static_cast<LocalNotify*>(notifyInfo[waitNotifyIndex])->Wait(waitStream, dispatcher_));

        // 从流2通知主流 lookup流水 done，notify array中倒数第二个notify用于HdcsCollLookupResetUniqueHandle算子执行后的同步
        s32 recordNotifyIndex = waitNotifyIndex;
        Stream postStream(subStreamInfo[PIPELINE_RECOVER_VALUE_SUB_STREAM_INDEX]);
        CHK_RET(static_cast<LocalNotify*>(notifyInfo[recordNotifyIndex])->Post(postStream, dispatcher_));
    }

    CHK_RET(profilingManager.CallEsMsprofReportTaskApi(false, beginTime, ProfTaskType::TASK_RESET_UNIQUE_HANDLE));
    return HCCL_SUCCESS;
}

HcclResult OpRemoteLookUpExecutor::WaitSendKeyFinish(OpTaskPtr &task) const
{
    // HdcsCollLookupWaitSendKeyFinish和HdcsLookupPipelineParams需要的参数一致，这里复用task结构体
    ResetUniqueHandleTask *taskPtr = dynamic_cast<ResetUniqueHandleTask *>(task.get());
    std::vector<rtStream_t> &subStreamInfo = taskPtr->subStreamInfo;
    std::vector<HcclRtNotify> &notifyInfo = taskPtr->notifyInfo;

    HdcsLookupPipelineParams para{};
    para.tag = taskPtr->tag;
    para.cubeIndex = taskPtr->cubeIndex;
    para.usePipeline = taskPtr->usePipeline;

    rtKernelLaunchNames_t launchNames;
    launchNames.kernelName = "HdcsCollLookupWaitSendKeyFinish";
    launchNames.soName = "libccl_kernel.so";
    launchNames.opName = "CollLookupWaitSendKeyFinish";
    CHK_RET(hrtAicpuKernelLaunch(launchNames, 1, &para, sizeof(HdcsLookupPipelineParams), nullptr,
        subStreamInfo[PIPELINE_SEND_KEYS_SUB_STREAM_INDEX]));

    // 等待 lookup send key done，notify array中倒数第一个notify用于HdcsCollLookupWaitSendKeyFinish算子执行后的同步
    s32 waitNotifyIndex = notifyInfo.size() - 1;
    Stream waitStream(taskPtr->stream);
    CHK_RET(static_cast<LocalNotify*>(notifyInfo[waitNotifyIndex])->Wait(waitStream, dispatcher_));

    // 从流0通知主流 lookup send key done
    // notify array中倒数第一个notify用于HdcsCollLookupWaitSendKeyFinish算子执行后的同步
    s32 recordNotifyIndex = notifyInfo.size() - 1;
    Stream postStream(subStreamInfo[PIPELINE_SEND_KEYS_SUB_STREAM_INDEX]);
    CHK_RET(static_cast<LocalNotify*>(notifyInfo[recordNotifyIndex])->Post(postStream, dispatcher_));
    return HCCL_SUCCESS;
}

TransportHandle OpRemoteLookUpExecutor::GetTransport(TransTask *taskPtr)
{
    TransportHandle transPtr = nullptr;
    if (taskPtr->dstRank == INVALID_VALUE_RANKID) {
        if (tempTransQue_.size() != 0) {
            transPtr = tempTransQue_.front();
            tempTransQue_.pop();
        }
    } else {
        transPtr = rankTransportMap_[taskPtr->dstRank];
    }
    return transPtr;
}

HcclResult OpRemoteLookUpExecutor::TaskIsend(OpTaskPtr &task)
{
    uint64_t beginTime = hrtMsprofSysCycleTime();
    TransTask *taskPtr = dynamic_cast<TransTask *>(task.get());
    TransportHeterog *heterogPtr = static_cast<TransportHeterog *>(GetTransport(taskPtr));
    CHK_PTR_NULL(heterogPtr);
    HcclRequestInfo *request = nullptr;

    if (taskPtr->count != 0) {
        CHK_RET(Isend(taskPtr, heterogPtr, request));
        request->tag = taskPtr->tag;
        HCCL_DEBUG("isend success tag %d request %p", taskPtr->tag, request);
        if (taskPtr->outRequest != nullptr) {
            *taskPtr->outRequest = request;
        }
    } else {
        // worker 侧没有需要查询的key, 需要record让worker侧流程走下去
        TransportHeterogRoce *transportHeterogRoce = dynamic_cast<TransportHeterogRoce *>(heterogPtr);
        // ES拉远推理是TransportRoce
        TransportRoce *transportRoce = dynamic_cast<TransportRoce *>(heterogPtr);

        // ES训练的RDMA直接发送notify
        if ((transportHeterogRoce != nullptr) && (transportRoce == nullptr)) {
            Stream tmpStream(nullptr);
            CHK_RET(transportHeterogRoce->RecordNotifyWithReq(tmpStream, RdmaNotifyOp::SEND_NOTIFY, request));
        } else {
            TransportEndPointInfo srcEp(0, taskPtr->srcRank, taskPtr->tag);
            TransportEndPointInfo dstEp(0, taskPtr->dstRank, taskPtr->tag);
            TransportEndPointParam epParam(srcEp, dstEp);
            TransData sendData(reinterpret_cast<u64>(nullptr), reinterpret_cast<u64>(nullptr), taskPtr->count,
                taskPtr->dataType);

            CHK_RET(heterogPtr->Send(sendData, epParam));
        }

        HCCL_DEBUG("empty envelop need send record to worker!!!");
    }

    auto &profilingManager = ProfilingManager::Instance();
    CHK_RET(profilingManager.CallEsMsprofReportTaskApi(false, beginTime, ProfTaskType::TASK_LOOKUP_RESPONSE_ISEND));

    return HCCL_SUCCESS;
}

HcclResult OpRemoteLookUpExecutor::TaskWait(OpTaskPtr &task)
{
    HCCL_INFO("Task Wait begin");
    WaitSomeTask *taskPtr = dynamic_cast<WaitSomeTask *>(task.get());
    *taskPtr->compCount = 0;
    s32 flag = HCCL_TEST_INCOMPLETED;
    // 有一个完成的request就退出循环
    while (flag != HCCL_TEST_COMPLETED) {
        if (!blockFlag_) {
            HCCL_INFO("TaskWait execute abort self");
            return HCCL_E_AGAIN;
        }
        for (s32 i = 0; i < taskPtr->requestCount; i++) {
            HcclRequestInfo *request = static_cast<HcclRequestInfo *>(taskPtr->requestArray[i]);
            TransportHeterog *transportPtr = static_cast<TransportHeterog *>(request->transportHandle);
            HcclStatus compState;
            CHK_RET(transportPtr->Test(*request, flag, compState));
            if (flag == HCCL_TEST_COMPLETED) {
                taskPtr->compIndices[*taskPtr->compCount] = i;
                taskPtr->compStatus[*taskPtr->compCount] = compState;
                (*taskPtr->compCount)++;
                break;
            }
        }
    }
    HCCL_INFO("Task Wait success");
    return HCCL_SUCCESS;
}
} // namespace hccl
